"""
created: 2023-03-26 14:05
author:qinguoming
description: 
1.  这个程序是一个计算模板,用于对3D线性8节点单元组成的结构的静态分析,输出节点位移,输出单元应力,单元应变,节点支反力等结果.
2.  节点自由度为: ux,uy,uz      
3.  节点和网格的编号必须从1开始,且连续.
4.  所有set_name,surface_name必须是小写
5.  计算结果输出到mat文件
"""
import sys
import os
import dill
from EF2D import *
import numpy as np
import scipy as sc
from icecream import ic
from EF2D.tools import *
from EF2D.tools import *
from EF2D.C3D8 import C3D8
from EF2D.Constants import *
from typing import List,Dict
from datetime import datetime
import matplotlib.pyplot as plt
initNumpyPrintOptions()
# -------------定义变量------------------
job_name="job-c3d8-case250407-1"
# 网格文件
# mesh_file="L:\\EF2D\\tests\\data\\Job-F300.inp"
mesh_file="L:/EF2D/tests/AbaqusModelTests/250407/Job250407C3D8COMPLEXMODEL.inp"
# 材料属性(SI-mm单位制)
thermal_expand= 14.4
rho= 1.932e-05
E=77200
nu=0.3
# 约束:支持位移边界条件
def_bc=[("DisplacementBc","Set_fix",1,0.0),
        ("DisplacementBc","Set_fix",2,0.0),
        ("DisplacementBc","Set_fix",3,0.0)]
# 定义载荷(载荷类型,载荷参数)
# def_load=[("ConcentratedLoad","load-1","Set_forceX_100N",1,300.0),]
def_load=[("ConcentratedLoad","load-1","Set_Fy100N",2,-100.0),]
# def_load=[("BodyLoad" ,"body_force-1","body_force_gravity_nz",'z',-500.0)]
# def_load=[ ("SurfaceLoad" , "surface_load-1000pa" , "surface_load_tensile_1000N", [-1.0,0,0] ,1000.0)]
# def_load=[("PressureLoad" , "pressure_500pa" , "pressure_load_500mpa", 5.0)]
# def_load=[ ("ThermalLoad" , "x" , "elset-1＂,10000.0)]

#-////////求解设置//////////////////////
bcMethod=2  # 边界条件处理方法,为1,2,3,4分别是Solver.py的方法
export_result=True  # 是否输出结果
export_path="L:/EF2D/tests/result/"
preHDF5File="C3D8_Static_pre.h5"
SolvedHDF5File="C3D8_Static_solved.h5"
postHDF5File="C3D8_Static_post.h5"

# -------------导入网格------------------
node_data, element_data, nset_data, elset_data ,surf_data= read_abaqus_inp(mesh_file)
print("网格文件:",mesh_file)
print("节点数:",node_data.shape[0])
print("单元数:",element_data.shape[0])
print("节点集:",nset_data.keys())
print("单元集:",elset_data.keys())
print("surface集合:",surf_data.keys())
print("边界条件:",[b[0] for b in def_bc])
print("载荷有:",[l[1] for l in def_load])
print("边界条件处理方法:",bcMethod)

# -------------建立网格------------------
node_num=node_data.shape[0]
element_num=element_data.shape[0]
flobal_dof_num=node_num*3
# 定义节点对象
node_dict:Dict[int,Node]={}
for i in range(node_num):
    if node_data[i,0] not in node_dict.keys():
        node_dict[node_data[i,0]]=Node(label=node_data[i,0],x=node_data[i,1],y=node_data[i,2],z=node_data[i,3])
# 定义单元对象
element_list:List[C3D8]=[]
for i in range(element_num):
    element_list.append(C3D8(label=element_data[i,0],
                                        node1=node_dict[element_data[i,1]],
                                        node2=node_dict[element_data[i,2]],
                                        node3=node_dict[element_data[i,3]],
                                        node4=node_dict[element_data[i,4]],
                                        node5=node_dict[element_data[i,5]],
                                        node6=node_dict[element_data[i,6]],
                                        node7=node_dict[element_data[i,7]],
                                        node8=node_dict[element_data[i,8]],
                                        E=E,nu=nu,density=rho,expand=thermal_expand))
# 定义载荷和边界
bc_list=[]
for b in def_bc:
    if b[0]=="DisplacementBc":
        if isinstance(b[1],str) and b[1].lower() in nset_data.keys():
            bc_list.append(DisplacementBc(name=b[0],nodeLabels=nset_data[b[1].lower()],dofth=b[2],value=b[3]))
        if (isinstance(b[1],list) or isinstance(b[1],np.ndarray)):
            bc_list.append(DisplacementBc(*b))
load_list=[]
for l in def_load:
    if l[0]=="ConcentratedLoad":
        if isinstance(l[2],str) and l[2].lower() in nset_data.keys():
            load_list.append(ConcentratedLoad(name=l[1],
                                            nodeLabels=nset_data[l[2].lower()],
                                            dofth=l[3],
                                            value=l[4]))
        else:
            load_list.append(ConcentratedLoad(*l[1:]))
    elif l[0]=="PressureLoad":
        if isinstance(l[2],str) and l[2].lower() in surf_data.keys():
            load_list.append(PressureLoad(name=l[1],
                                        ElementLabels=surf_data[l[2].lower()][:,0],
                                        faceIds=surf_data[l[2].lower()][:,1],
                                        value=l[3]))
        else:
            load_list.append(PressureLoad(name=l[1],
                                        ElementLabels=l[2][:,0],
                                        faceIds=l[2][:,1],
                                        value=l[3]))
    elif l[0] == "BodyLoad":
        if isinstance(l[2],str) and l[2].lower() in elset_data.keys():
            load_list.append(BodyLoad(name=l[1],
                                    ElementLabels=elset_data[l[2].lower()],
                                    Direction=l[3],
                                    value=l[4]))
        else:
            load_list.append(BodyLoad(*l[1:]))
    elif l[0] == "SurfaceLoad":
        if isinstance(l[2],str) and l[2].lower() in surf_data.keys():
            load_list.append(SurfaceLoad(name=l[1],
                                        ElementLabels=surf_data[l[2].lower()][:,0],
                                        faceIds=surf_data[l[2].lower()][:,1],
                                        directionVector=l[3],
                                        value=l[4]))
        else:
            load_list.append(SurfaceLoad(name=l[1],
                                        ElementLabels=l[2][:,0],
                                        faceIds=l[2][:,1],
                                        directionVector=l[3],
                                        value=l[4]))
    # ("ThermalLoad" , "y" ,element_label_list,  5.0)
    elif l[0] == "ThermalLoad":
        if isinstance(l[2],str) and l[2].lower() in elset_data.keys():
            load_list.append(ThermalLoad(name=l[1],
                                    ElementLabels=elset_data[l[2].lower()],
                                    temperature=l[4]))
        else:
            load_list.append(BodyLoad(*l[1:]))
# -------------组装刚度矩阵-------------------------
Kg=np.zeros((flobal_dof_num,flobal_dof_num))
for e in element_list:
    Kg[np.ix_(e.ElemDofIndexs,e.ElemDofIndexs)]+=e.Ke_bbar
# --------------边界条件处理-------------------------
bc_dict={}
for bc in bc_list:
    if isinstance(bc,DisplacementBc):
        for n_label in bc.NodeLabelList:
            dofind=get_ithdof_index(n_label,3,bc.Dofth)
            if dofind not in bc_dict.keys():
                bc_dict[dofind]=bc.Value
            else:
                bc_dict[dofind]+=bc.Value
# key: 从1开始的自由度编号, value: 加载节点力值
eqf=np.zeros(flobal_dof_num)
for load in load_list:
    if isinstance(load,ConcentratedLoad):
        for n_label in load.get_NodesList:
            dof=get_ithdof_index(n_label,3,load.dofth)
            eqf[dof]+=load.LoadValue
    else:
        # 处理单元载荷
        for e in element_list:
            if e.label not in load.get_ElemsList:
                continue
            eqf[e.ElemDofIndexs]+=e.ProcessELoad2EqNodeForce(load)
# 找出eqF中非0值的索引和非零值,定义全局节点力字典{key:自由度索引(0-base), value:节点力}
# P_ind=eqf.nonzero()
# load_dict=dict(zip((P_ind[0]+1).tolist(),Fg[P_ind].tolist()))
load_dict={ind:eqf[ind] for ind in eqf.nonzero()[0] }

# --------------引入边界条件------------------

Kg_bc=Kg.copy()
dofNum=Kg.shape[0]
# 初始化向量
U,P_bc=np.zeros(dofNum),np.zeros(dofNum)
# 已知节点力加到 P
for ind,Pval in load_dict.items():
    P_bc[ind]+=Pval
# 循环所有的位移约束
for j,Uval in bc_dict.items():
    #Step1
    for i in range(dofNum):
        P_bc[i]=P_bc[i]-Kg_bc[i,j]*Uval
    #Step2
    for k in range(dofNum):
        Kg_bc[k,j]=0
        Kg_bc[j,k]=0
    Kg_bc[j,j]=1
    #Step3
    P_bc[j]=Uval
# 将 K'U=P'写入hdf5文件,给matlab求解
import h5py
with h5py.File("C3D8_Static_pre.h5","w") as f:
    f.create_dataset("Kg",data=Kg)
    f.create_dataset("Kg_bc",data=Kg_bc)
    f.create_dataset("F_bc",data=P_bc)

